from numpy.lib.stride_tricks import as_strided
import numpy as np
import pandas as pd

def fib_yield_for(n):
    a, b = 0, 1
    for _ in range(n):
        a, b = b, a + b
        yield a


def myroll(df, w, **kwargs):
    v = df.values
    d0, d1 = v.shape
    s0, s1 = v.strides
    a = as_strided(v, (d0 - (w - 1), w, d1), (s0, s0, s1))
    rolled_df = pd.concat({
        row: pd.DataFrame(values, columns=df.columns)
        for row, values in zip(df.index, a)
    })
    return rolled_df.groupby(level=0, **kwargs)


def roll_corr(x):
    print(x)
    r = x["close"].astype("float32").corr(x["szss_close"].astype("float32"))
    print(r, type(r))
    if np.isnan(r):
        return 0
    else:
        return r


## 速度太慢了
def myroll_apply_corr(df, w):
    """
     速度太慢， 放弃ing
    :param df:
    :param w:
    :return:
    """
    df["corr" + str(w)] = myroll(df, w).apply(roll_corr)
    df["corr" + str(w)] = df["corr" + str(w)].shift(w-1)
    return df


def corr_v1(df1, col1, col2, n_list=None):
    if n_list is None:
        n_list = [5, 10, 15, 20]
    for n in n_list:
        df1["corr" + str(n)] = df1[col1].astype("float32").rolling(n).corr(df1[col2].astype("float32"))
    return df1


def corr_v2(df1, col1, col2, n_list=None):
    if n_list is None:
        n_list = [5, 10, 15, 20]
    for n in n_list:
        df1["corr" + str(n)] = pd.rolling_corr(df1[col1].astype("float32"), df1[col2].astype("float32"), n)
    return df1


if __name__ == '__main__':
    j = 0
    for i in fib_yield_for(10):
        j = j + 1
        print(j, i)
